package com.RJAI.gemini.controller;


import cn.hutool.core.util.ObjectUtil;
import cn.hutool.core.util.RandomUtil;
import cn.hutool.core.util.StrUtil;
import cn.hutool.http.HttpRequest;
import cn.hutool.http.HttpResponse;
import cn.hutool.http.HttpUtil;
import cn.hutool.json.JSONObject;
import cn.hutool.json.JSONUtil;
import com.RJAI.gemini.utils.GCloudUtils;
import com.RJAI.gemini.exception.RJException;

import com.RJAI.gemini.service.impl.AsyncService;
import okhttp3.*;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.data.redis.core.RedisTemplate;
import org.springframework.http.MediaType;
import org.springframework.retry.annotation.Backoff;
import org.springframework.retry.annotation.Retryable;
import org.springframework.web.bind.annotation.*;
import org.springframework.web.bind.annotation.RequestBody;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;

import javax.annotation.Resource;
import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;
import java.io.IOException;
import java.util.Arrays;
import java.util.HashMap;
import java.util.List;
import java.util.concurrent.TimeUnit;
/*
作者信息
WeChat:ciansec
QQ:1068092013
 */
@RestController
public class GeminiAiController {


    @Autowired
    private AsyncService asyncService;

    @Resource
    private RedisTemplate redisTemplate;

    @Value("${cian.serviceAccountDir}")
    private String serviceAccountDir;

    @Value("${cian.serviceAccountExpiration}")
    private Integer serviceAccountExpiration;


    //代理GeminiAI非流式对话
    @Retryable(maxAttempts = 10, backoff = @Backoff(delay = 2000))
    @PostMapping({"/{key}/v1/models/{modelVersion}:generateContent", "/{key}/v1beta/models/{modelVersion}:generateContent"})
    public JSONObject getGeminiAi(@PathVariable("modelVersion") String modelVersion, @PathVariable("key") String key, @RequestBody JSONObject jsonObject) throws InterruptedException, IOException {
        //用户没有传入模型名称时默认的模型名称
        if (StrUtil.isEmpty(modelVersion)) {
            modelVersion = "gemini-1.5-pro";
        }



        //todo 这里地区分的不够细致 每个区域可用模型不一样
        List<String> regions = Arrays.asList("asia-east1","asia-east2","asia-northeast1","asia-northeast3","asia-south1","asia-southeast1","australia-southeast1","europe-central2","europe-north1","europe-southwest1","europe-west1","europe-west2","europe-west3","europe-west4","europe-west6"," europe-west8"," europe-west9","me-central1","me-central2","me-west1","northamerica-northeast1","southamerica-east1","us-central1","us-east1","us-east4","us-east5","us-south1","us-west1","us-west4");
        //视觉模型
        if(modelVersion.contains("vision")){
            regions = Arrays.asList("asia-east1","asia-east2","asia-northeast1","asia-northeast3","asia-south1","asia-southeast1","australia-southeast1","europe-central2","europe-north1","europe-southwest1","europe-west1","europe-west2","europe-west3","europe-west4","europe-west6"," europe-west8"," europe-west9","me-central1","me-central2","me-west1","northamerica-northeast1","southamerica-east1","us-central1","us-east1","us-east4","us-east5","us-south1","us-west1","us-west4");
        }else if (modelVersion.contains("gemini-pro") || modelVersion.contains("gemini-1.0-pro")){
            regions = Arrays.asList("asia-east1","asia-east2","asia-northeast1","asia-northeast3","asia-south1","asia-southeast1","australia-southeast1","europe-central2","europe-north1","europe-southwest1","europe-west1","europe-west2","europe-west3","europe-west4","europe-west6"," europe-west8"," europe-west9","me-central1","me-central2","me-west1","northamerica-northeast1","southamerica-east1","us-central1","us-east1","us-east4","us-east5","us-south1","us-west1","us-west4");
        }else if (modelVersion.contains("gemini-1.5-flash")){
            regions = Arrays.asList("asia-east1","asia-east2","asia-northeast1","asia-northeast3","asia-south1","asia-southeast1","australia-southeast1","europe-central2","europe-north1","europe-southwest1","europe-west1","europe-west2","europe-west3","europe-west4","europe-west6"," europe-west8"," europe-west9","me-central1","me-central2","me-west1","northamerica-northeast1","southamerica-east1","us-central1","us-east1","us-east4","us-east5","us-south1","us-west1","us-west4");
        }


        String region = regions.get(RandomUtil.randomInt(0, regions.size() - 1)).replace(" ", "");

        //获得Gcloud的token和projectID
        String accessToken = null;
        String projectId = null;
        //全局锁
        synchronized ("getkey") {
            String currentKey = key;
            System.out.println("current_key = " + currentKey);
            //判断Redis中是否存在对应key的信息，如果有直接用
            if (redisTemplate.hasKey(currentKey + "::accessToken")) {
                accessToken = (String) redisTemplate.opsForValue().get(currentKey + "::accessToken");
                projectId = (String) redisTemplate.opsForValue().get(currentKey + "::project_id");
            } else {
                //如果没有就去生成
                HashMap<String, String> map = GCloudUtils.generateGCloudJWTByJson(serviceAccountDir + "/"+currentKey.replace( "|","/" ) + ".json");
                accessToken = map.get("accessToken");
                projectId = map.get("project_id");
                redisTemplate.opsForValue().set(currentKey + "::accessToken", accessToken, serviceAccountExpiration, TimeUnit.MINUTES);
                redisTemplate.opsForValue().set(currentKey + "::project_id", projectId, serviceAccountExpiration, TimeUnit.MINUTES);
            }
        }

        String url = "https://" + region + "-aiplatform.googleapis.com/v1/projects/" + projectId + "/locations/" + region + "/publishers/google/models/" + modelVersion + ":generateContent";
        HttpRequest request = HttpUtil.createPost(url).header("Content-Type", "application/json").header("Authorization", "Bearer " + accessToken).body(jsonObject.toString());
        request.setConnectionTimeout(999999);
        request.setReadTimeout(999999);
        HttpResponse execute = null;
        try {
            execute = request.execute();

        } catch (Exception e) {
            e.printStackTrace();
            throw new RJException("请求Gemini服务器失败", -1);

        }


        //如果返回数据不是空的
        if (ObjectUtil.isNotEmpty(execute)) {
            String tempContent = "";
            tempContent = execute.body();
            System.out.println("tempContent = " + tempContent);
            if (StrUtil.isNotBlank(tempContent)) {
                try {
                    JSONObject object = JSONUtil.parseObj(tempContent);

                    try {
                        if (object.getJSONObject("error").getInt("code") == 429) {
                            throw new RJException("429错误。", -1);
                        }
                    } catch (Exception e) {
                        System.out.println("e1 = " + e);
                    }
                    //这里其实没写完 应该要判断请求来的max_tokens是不是1，因为1是oneai测试的数量
                    try {
                        if ("MAX_TOKENS".equals(object.getJSONArray("candidates").getJSONObject(0).getStr("finishReason"))) {
                            return object;
                        }
                    } catch (Exception e) {
                        System.out.println("e2 = " + e);
                    }

                    String text = object.getJSONArray("candidates").getJSONObject(0).getJSONObject("content").getJSONArray("parts").getJSONObject(0).getStr("text");
                    if (StrUtil.isNotBlank(text)) {
                        return object;
                    } else {
                        return object;
                    }
                } catch (Exception e) {
                    throw new RJException("数据解析异常。", -1);
                }
            } else {

                throw new RJException("Gemini返回失败2", -1);
            }
        } else {
            throw new RJException("Gemini内容为空", -1);
        }


    }



    //代理GeminiAI流式对话
    @PostMapping(value = {"/{key}/v1/models/{modelVersion}:streamGenerateContent", "/{key}/v1beta/models/{modelVersion}:streamGenerateContent"}, produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    public SseEmitter generateContentStream(HttpServletResponse response, @PathVariable("modelVersion") String modelVersion, @PathVariable("key") String key, @RequestBody JSONObject jsonObject) throws InterruptedException, IOException {
        /*response.addHeader( "Transfer-Encoding","chunked" );*/
        response.addHeader("X-Content-Type-Options", "nosniff");
        response.addHeader("X-Frame-Options", "SAMEORIGIN");
        response.addHeader("X-XSS-Protection", "no");
        response.addHeader("Content-Disposition", "attachment");

        SseEmitter sseEmitter = new SseEmitter(0L);


        String accessToken = null;
        String projectId = null;
        synchronized (key) {
            String current_key = key;
            System.out.println("current_key = " + current_key);
            if (redisTemplate.hasKey(current_key + "::accessToken")) {
                accessToken = (String) redisTemplate.opsForValue().get(current_key + "::accessToken");
                projectId = (String) redisTemplate.opsForValue().get(current_key + "::project_id");
            } else {
                HashMap<String, String> map = GCloudUtils.generateGCloudJWTByJson(serviceAccountDir +"/"+ current_key.replace( "|","/" ) + ".json");
                accessToken = map.get("accessToken");
                projectId = map.get("project_id");
                redisTemplate.opsForValue().set(current_key + "::accessToken", accessToken, serviceAccountExpiration, TimeUnit.MINUTES);
                redisTemplate.opsForValue().set(current_key + "::project_id", projectId, serviceAccountExpiration, TimeUnit.MINUTES);
            }
        }

        //todo 这里地区分的不够细致 每个区域可用模型不一样
        List<String> regions = Arrays.asList("asia-east1","asia-east2","asia-northeast1","asia-northeast3","asia-south1","asia-southeast1","australia-southeast1","europe-central2","europe-north1","europe-southwest1","europe-west1","europe-west2","europe-west3","europe-west4","europe-west6"," europe-west8"," europe-west9","me-central1","me-central2","me-west1","northamerica-northeast1","southamerica-east1","us-central1","us-east1","us-east4","us-east5","us-south1","us-west1","us-west4");
        //视觉模型
        if(modelVersion.contains("vision")){
            regions = Arrays.asList("asia-east1","asia-east2","asia-northeast1","asia-northeast3","asia-south1","asia-southeast1","australia-southeast1","europe-central2","europe-north1","europe-southwest1","europe-west1","europe-west2","europe-west3","europe-west4","europe-west6"," europe-west8"," europe-west9","me-central1","me-central2","me-west1","northamerica-northeast1","southamerica-east1","us-central1","us-east1","us-east4","us-east5","us-south1","us-west1","us-west4");
        }else if (modelVersion.contains("gemini-pro") || modelVersion.contains("gemini-1.0-pro")){
            regions = Arrays.asList("asia-east1","asia-east2","asia-northeast1","asia-northeast3","asia-south1","asia-southeast1","australia-southeast1","europe-central2","europe-north1","europe-southwest1","europe-west1","europe-west2","europe-west3","europe-west4","europe-west6"," europe-west8"," europe-west9","me-central1","me-central2","me-west1","northamerica-northeast1","southamerica-east1","us-central1","us-east1","us-east4","us-east5","us-south1","us-west1","us-west4");
        }else if (modelVersion.contains("gemini-1.5-flash")){
            regions = Arrays.asList("asia-east1","asia-east2","asia-northeast1","asia-northeast3","asia-south1","asia-southeast1","australia-southeast1","europe-central2","europe-north1","europe-southwest1","europe-west1","europe-west2","europe-west3","europe-west4","europe-west6"," europe-west8"," europe-west9","me-central1","me-central2","me-west1","northamerica-northeast1","southamerica-east1","us-central1","us-east1","us-east4","us-east5","us-south1","us-west1","us-west4");
        }


        String region = regions.get(RandomUtil.randomInt(0, regions.size() - 1)).replace(" ", "");
        String url = "https://" + region + "-aiplatform.googleapis.com/v1/projects/" + projectId + "/locations/" + region + "/publishers/google/models/" + modelVersion + ":streamGenerateContent?alt=sse";
        OkHttpClient.Builder clientBuilder = new OkHttpClient.Builder();
        okhttp3.RequestBody body = okhttp3.RequestBody.create(okhttp3.MediaType.get("application/json; charset=utf-8"), jsonObject.toString());
        Request request = new Request.Builder()
                .url(url)
                .post(body)
                .addHeader("Content-Type", "application/json")
                .addHeader("Authorization", "Bearer " + accessToken)
                .build();
        System.out.println("---------------------流式传输开始---------------------");
        System.out.println("request = " + request);
        asyncService.asyncGemini(clientBuilder, request, sseEmitter);
        System.out.println("---------------------流式传输异步OK---------------------");
        return sseEmitter;
    }

}
